"""《RRT*算法实现》
    时间：2025.06.25
"""
import math
import os
import sys

import matplotlib.pyplot as plt
from celluloid import Camera  # 保存动图时用，pip install celluloid

sys.path.append("../RRT")
try:
	from rrt_planning import RRT
except ImportError:
	raise

show_animation = True


class RRTStar(RRT):
	"""
	Class for RRT Star planning
	"""

	class Node(RRT.Node):
		def __init__(self, x, y):
			super().__init__(x, y)
			self.cost = 0.0

	def __init__(self,
				 start,
				 goal,
				 obstacle_list,
				 rand_area,
				 expand_dis=3.0,
				 goal_sample_rate=20,
				 max_iter=50000,
				 connect_circle_dist=50.0,
				 search_until_max_iter=False,
				 robot_radius=0.0):
		"""
		Setting Parameter
		start:Start Position [x,y]
		goal:Goal Position [x,y]
		obstacleList:obstacle Positions [[x,y,size],...]
		randArea:Random Sampling Area [min,max]
		"""
		super().__init__(start, goal, obstacle_list, rand_area, expand_dis,
						 goal_sample_rate, max_iter,
						 robot_radius=robot_radius)
		self.connect_circle_dist = connect_circle_dist
		self.goal_node = self.Node(goal[0], goal[1])
		self.search_until_max_iter = search_until_max_iter

	def planning(self, animation=True, camera=None):
		"""
		rrt star path planning
		animation: flag for animation on or off .
		"""

		self.node_list = [self.start]
		for i in range(self.max_iter):
			if i % 100 == 0:
				print("Iter:", i, ", number of nodes:", len(self.node_list))
			# print("Iter:", i, ", number of nodes:", len(self.node_list))
			rnd = self.sample_free()
			nearest_ind = self.get_nearest_node_index(self.node_list, rnd)
			new_node = self.steer(self.node_list[nearest_ind], rnd,
								  self.expand_dis)
			near_node = self.node_list[nearest_ind]
			# 计算代价,欧氏距离
			new_node.cost = near_node.cost + math.hypot(new_node.x - near_node.x, new_node.y - near_node.y)

			if self.obstacle_free(new_node, self.obstacle_list, self.robot_radius):
				near_inds = self.find_near_nodes(new_node)  # 找到x_new的邻近节点
				node_with_updated_parent = self.choose_parent(new_node, near_inds)  # 重新选择父节点
				# 如果父节点更新了
				if node_with_updated_parent:
					# 重布线
					self.rewire(node_with_updated_parent, near_inds)
					self.node_list.append(node_with_updated_parent)
				else:
					self.node_list.append(new_node)

			if animation and i % 100 == 0:
				self.draw_graph(rnd, camera)

			if ((not self.search_until_max_iter) and new_node):  # if reaches goal
				last_index = self.search_best_goal_node()
				if last_index is not None:
					return self.generate_final_course(last_index)

		print("reached max iteration")

		last_index = self.search_best_goal_node()
		if last_index is not None:
			return self.generate_final_course(last_index)

		return None

	def choose_parent(self, new_node, near_inds):
		"""
		在新产生的节点 $x_{new}$ 附近以定义的半径范围$r$内寻找所有的近邻节点 $X_{near}$，
		作为替换 $x_{new}$ 原始父节点 $x_{near}$ 的备选
		我们需要依次计算起点到每个近邻节点 $X_{near}$ 的路径代价 加上近邻节点 $X_{near}$ 到 $x_{new}$ 的路径代价，
		取路径代价最小的近邻节点$x_{min}$作为 $x_{new}$ 新的父节点
		Arguments:
		--------
			new_node, Node
				randomly generated node with a path from its neared point
				There are not coalitions between this node and th tree.
			near_inds: list
				Indices of indices of the nodes what are near to new_node
		Returns.
		------
			Node, a copy of new_node
		"""
		if not near_inds:
			return None

		# search nearest cost in near_inds
		costs = []
		for i in near_inds:
			near_node = self.node_list[i]
			t_node = self.steer(near_node, new_node)
			if t_node and self.obstacle_free(t_node, self.obstacle_list, self.robot_radius):
				costs.append(self.calc_new_cost(near_node, new_node))
			else:
				costs.append(float("inf"))  # the cost of collision node
		min_cost = min(costs)

		if min_cost == float("inf"):
			print("There is no good path.(min_cost is inf)")
			return None

		min_ind = near_inds[costs.index(min_cost)]
		new_node = self.steer(self.node_list[min_ind], new_node)
		new_node.cost = min_cost

		return new_node

	def search_best_goal_node(self):
		dist_to_goal_list = [self.calc_dist_to_goal(n.x, n.y) for n in self.node_list]
		goal_inds = [
			dist_to_goal_list.index(i) for i in dist_to_goal_list
			if i <= self.expand_dis
		]

		safe_goal_inds = []
		for goal_ind in goal_inds:
			t_node = self.steer(self.node_list[goal_ind], self.goal_node)
			if self.obstacle_free(t_node, self.obstacle_list, self.robot_radius):
				safe_goal_inds.append(goal_ind)

		if not safe_goal_inds:
			return None

		min_cost = min([self.node_list[i].cost for i in safe_goal_inds])
		for i in safe_goal_inds:
			if self.node_list[i].cost == min_cost:
				return i

		return None

	def find_near_nodes(self, new_node):
		"""
		1) defines a ball centered on new_node
		2) Returns all nodes of the three that are inside this ball
			Arguments:
			---------
				new_node: Node
					new randomly generated node, without collisions between
					its nearest node
			Returns:
			-------
				list
					List with the indices of the nodes inside the ball of
					radius r
		"""
		nnode = len(self.node_list) + 1
		r = self.connect_circle_dist * math.sqrt((math.log(nnode) / nnode))
		# if expand_dist exists, search vertices in a range no more than
		# expand_dist
		if hasattr(self, 'expand_dis'):
			r = min(r, self.expand_dis)
		dist_list = [(node.x - new_node.x) ** 2 + (node.y - new_node.y) ** 2
					 for node in self.node_list]
		near_inds = [dist_list.index(i) for i in dist_list if i <= r ** 2]
		return near_inds

	def rewire(self, new_node, near_inds):
		"""
			For each node in near_inds, this will check if it is cheaper to
			arrive to them from new_node.
			In such a case, this will re-assign the parent of the nodes in
			near_inds to new_node.
			Parameters:
			----------
				new_node, Node
					Node randomly added which can be joined to the tree
				near_inds, list of uints
					A list of indices of the self.new_node which contains
					nodes within a circle of a given radius.
			Remark: parent is designated in choose_parent.
		"""
		for i in near_inds:
			near_node = self.node_list[i]
			edge_node = self.steer(new_node, near_node)
			if not edge_node:
				continue
			edge_node.cost = self.calc_new_cost(new_node, near_node)

			no_collision = self.obstacle_free(
				edge_node, self.obstacle_list, self.robot_radius)
			improved_cost = near_node.cost > edge_node.cost

			if no_collision and improved_cost:
				near_node.x = edge_node.x
				near_node.y = edge_node.y
				near_node.cost = edge_node.cost
				near_node.path_x = edge_node.path_x
				near_node.path_y = edge_node.path_y
				near_node.parent = edge_node.parent
				self.propagate_cost_to_leaves(new_node)

	def calc_new_cost(self, from_node, to_node):
		d, _ = self.calc_distance_and_angle(from_node, to_node)
		return from_node.cost + d

	def propagate_cost_to_leaves(self, parent_node):

		for node in self.node_list:
			if node.parent == parent_node:
				node.cost = self.calc_new_cost(parent_node, node)
				self.propagate_cost_to_leaves(node)


def main():
	print("Start ")
	fig = plt.figure(1)

	camera = Camera(fig)  # 保存动图时使用
	# camera = None # 不保存动图时，camara为None
	show_animation = True
	# ====Search Path with RRT====

	# # 生成100个随机元组的列表
	import random
	# obstacle_list = [(random.randint(0, 200), random.randint(0, 200), random.randint(1, 10)) for _ in range(100)]
	# 分别生成100个x、y坐标对和100个size值

	positions = [(random.randint(0, 180), random.randint(0, 180)) for _ in range(100)]
	sizes = [random.randint(1, 10) for _ in range(100)]

	# 合并为100个三元组
	obstacle_list = [(x, y, size) for (x, y), size in zip(positions, sizes)]

	while True:
		goal = [random.randint(180, 199), random.randint(180, 199)]
		goal_ = tuple(goal)
		if goal_ not in positions:
			break

	# Set Initial parameters
	rrt_star = RRTStar(
		start=[0, 0],
		goal=goal,
		rand_area=[0, 200],
		obstacle_list=obstacle_list,
		expand_dis=3,
		robot_radius=0.8)
	path = rrt_star.planning(animation=show_animation, camera=camera)

	if path is None:
		print("Cannot find path")
	else:
		print(f"found path!!,路径长度{len(path)},路径{path}")

		# Draw final path
		if show_animation:
			rrt_star.draw_graph(camera=camera)
			plt.plot([x for (x, y) in path], [y for (x, y) in path], 'r--')
			plt.grid(True)
			if camera != None:
				camera.snap()
				animation = camera.animate()
				animation.save('trajectory.gif')
	plt.show()


if __name__ == '__main__':
	main()